#include <cassert>
#include <chrono>
#include <iostream>
#include <mpi.h>
#include <random>
#include <vector>

std::random_device rd;

class Matrix {
private:
  int rows_;
  int cols_;

  std::mt19937 gen;

public:
  std::vector<double> data_;

  Matrix(int rows, int cols) : rows_(rows), cols_(cols) {
    data_.resize(rows * cols);
    this->gen = std::mt19937(rd());
  }

  void random() {
    for (int i = 0; i < rows_; i++) {
      for (int j = 0; j < cols_; j++) {
        (*this)(i, j) = std::uniform_real_distribution<double>(0, 1)(gen);
      }
    }
  }

  int rows() const { return rows_; }
  int cols() const { return cols_; }

  double &operator()(int i, int j) { return data_[i * cols_ + j]; }
  double operator()(int i, int j) const { return data_[i * cols_ + j]; }

  bool operator==(const Matrix &B) const {
    if (rows_ != B.rows_ || cols_ != B.cols_) {
      return false;
    }
    for (int i = 0; i < rows_; i++) {
      for (int j = 0; j < cols_; j++) {
        if ((*this)(i, j) != B(i, j)) {
          return false;
        }
      }
    }
    return true;
  }

  Matrix operator*(const Matrix &B) {
    Matrix C(rows_, B.cols_);
    for (int i = 0; i < rows_; i++) {
      for (int j = 0; j < B.cols_; j++) {
        C(i, j) = 0;
        for (int k = 0; k < cols_; k++) {
          C(i, j) += (*this)(i, k) * B(k, j);
        }
      }
    }
    return C;
  }

  void print() {
    for (int i = 0; i < rows_; i++) {
      for (int j = 0; j < cols_; j++) {
        std::cout << (*this)(i, j) << " ";
      }
      std::cout << std::endl;
    }
  }
};

Matrix mpiMatrixMul(const Matrix &A, const Matrix &B, int rank, int size) {
  int n = A.rows();
  int m = B.cols();
  int p = A.cols();

  int rows_per_proc = n / size;
  int extra_rows = n % size;

  int start_row = rank * rows_per_proc + std::min(rank, extra_rows);
  int end_row = start_row + rows_per_proc + (rank < extra_rows ? 1 : 0);
  int local_rows = end_row - start_row;

  Matrix C_partial(local_rows, m);

  for (int i = start_row; i < end_row; i++) {
    for (int j = 0; j < m; j++) {
      double sum = 0;
      for (int k = 0; k < p; k++) {
        sum += A(i, k) * B(k, j);
      }
      C_partial(i - start_row, j) = sum;
    }
  }

  Matrix C(n, m);
  std::vector<int> sendcounts(size);
  std::vector<int> displs(size);

  for (int i = 0; i < size; ++i) {
    int start_row_i = i * rows_per_proc + std::min(i, extra_rows);
    int end_row_i = start_row_i + rows_per_proc + (i < extra_rows ? 1 : 0);
    sendcounts[i] = (end_row_i - start_row_i) * m;
    displs[i] = start_row_i * m;
  }

  MPI_Allgatherv(C_partial.data_.data(), local_rows * m, MPI_DOUBLE,
                 C.data_.data(), sendcounts.data(), displs.data(), MPI_DOUBLE,
                 MPI_COMM_WORLD);
  return C;
}

int main(int argc, char **argv) {
  MPI_Init(&argc, &argv);

  int rank, size;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);

  if (argc < 2) {
    if (rank == 0) {
      std::cerr << "Usage: " << argv[0] << " <matrix_size>\n";
    }
    MPI_Finalize();
    return 1;
  }

  int n = std::stoi(argv[1]);

  Matrix A(n, n);
  Matrix B(n, n);

  if (rank == 0) {
    A.random();
    B.random();
  }

  MPI_Bcast(B.data_.data(), n * n, MPI_DOUBLE, 0, MPI_COMM_WORLD);
  MPI_Bcast(A.data_.data(), n * n, MPI_DOUBLE, 0, MPI_COMM_WORLD);

  if (rank == 0) {
    auto start = std::chrono::high_resolution_clock::now();
    Matrix C = mpiMatrixMul(A, B, rank, size);
    auto end = std::chrono::high_resolution_clock::now();
    auto duration =
        std::chrono::duration_cast<std::chrono::microseconds>(end - start)
            .count();

    start = std::chrono::high_resolution_clock::now();
    Matrix C_pthread = A * B;
    end = std::chrono::high_resolution_clock::now();
    std::cout << "Normal: "
              << std::chrono::duration_cast<std::chrono::microseconds>(end -
                                                                       start)
                     .count()
              << std::endl;
    std::cout << "Pthread: " << duration << std::endl;
  } else {
    Matrix C = mpiMatrixMul(A, B, rank, size);
  }

  if (rank == 0) {
    assert(C == A * B);
  }

  MPI_Finalize();
  return 0;
}
